# -*- coding: utf-8 -*-
"""
@Time    : 2024/9/6 13:40 
@Author  : ZhangShenao 
@File    : 9.LangGraph的流式响应.py 
@Desc    : LangGraph的流式响应

在LangChain中,stream流式响应以token为单位的,即每次响应都是一个token的数据
但是与LangChain不同,LangGraph的流式响应是以节点为单位的
也就是说,LangGraph的流式响应每次输出的都是一个节点的数据状态

在LangGraph中,流式响应有两种模式:
1. values: 总量模式,每次响应都返回当前节点执行后的最新Graph状态
2. updates: 增量模式,每次响应都返回当前节点对Graph状态的更新操作,格式为dict,key为节点名称,value为对状态的更新操作

在调用Graph的stream()方法时,可以通过stream_mode指定流式状态

"""
import os

import dotenv
from langchain_core.messages import HumanMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import MessagesState
from langgraph.prebuilt import create_react_agent

from dalle_image_tool import create_dalle_image_tool
from gaode_weather_tool import GaodeWeatherTool
from google_serper_tool import create_google_serper_tool

# 加载环境变量
dotenv.load_dotenv()

# 创建工具列表
tools = [GaodeWeatherTool(), create_dalle_image_tool(), create_google_serper_tool()]

# 创建LLM
llm = ChatOpenAI(model="gpt-4o-mini",
                 temperature=0,
                 openai_api_base=os.getenv("OPENAI_API_BASE"))

# 使用LangGraph的prebuild函数,创建预构建好的ReACT智能体
# 该函数封装了初始化Graph、创建Node和Edge和编译Graph等过程
# 最终返回一个编译好的CompiledGraph
# 它是一个Runnable可运行组件,可以直接调用
agent = create_react_agent(model=llm,  # 指定LLM模型
                           tools=tools,  # 指定工具列表
                           )

# 定义初始状态
init_state: MessagesState = {"messages": [HumanMessage(content="请帮我画一张英短蓝猫在沙发上睡觉的图片")]}

# 调用智能体,流式响应
# updates模式,每次响应都返回当前节点对Graph状态的更新
# 响应格式为dict,key为节点名称,value为对状态的更新操作
# for chunk in agent.stream(input=init_state, stream_mode="updates"):
#     print("=====================================================")
#     print(chunk)


# values模式,每次响应都返回当前节点执行后的最新Graph状态
# 可以获取状态的消息列表中的最新一条消息
for chunk in agent.stream(input=init_state, stream_mode="values"):
    # langchain_core.messages.base.BaseMessage.pretty_print()方法,可以友好地打印消息内容
    print(chunk["messages"][-1].pretty_print())
